import torch
import torch.nn as nn
import torch.nn.functional as F

class EncoderRNN(nn.Module):
    def __init__(self,input_size,hidden_size,dropout=0.1):
        super(EncoderRNN,self).__init__()
        self.hidden_size = hidden_size
        self.embedding = nn.Embedding(input_size,hidden_size)
        self.rnn = nn.RNN(hidden_size,hidden_size,batch_first=True)
        self.dropout = nn.Dropout(dropout)

    def forward(self,x):
        x = self.embedding(x)
        x = self.dropout(x)
        output,hidden = self.rnn(x)
        return output,hidden


encoder = EncoderRNN(input_size=10,hidden_size=5)
input_vector = torch.tensor([[0,1,2,3,4,5,6,7,8,9]])
SOS_token = torch.tensor(1)
output,hidden = encoder(input_vector)
print('输入向量维度：',input_vector.size())
print('输出向量维度：',output.size())
print('最终隐藏状态维度：',hidden.size())

class DecoderRNN(nn.Module):
    def __init__(self,hidden_size,output_size):
        super(DecoderRNN,self).__init__()
        self.embedding = nn.Embedding(output_size,hidden_size)
        self.rnn = nn.RNN(hidden_size,hidden_size,batch_first=True)
        self.out = nn.Linear(hidden_size,output_size)

    def forward(self,encoder_outputs,encoder_hidden,target_tensor=None):
        batch_size = encoder_outputs.size(0)
        decoder_input = torch.empty(batch_size,1,dtype=torch.long)
        decoder_hidden = encoder_hidden
        decoder_outputs = []
        for i in range(100):
            decoder_output,decoder_hidden = self.forward_step(decoder_input,decoder_hidden)
            decoder_outputs.append(decoder_output)
            if target_tensor is not None:
                decoder_input = target_tensor[:,i // 10].unsqueeze(1)
            else:
                _,topi = decoder_output.topk(1)
                decoder_input = topi.sequeeze(-1).detach()
        decoder_outputs = torch.cat(decoder_outputs,dim=1)
        decoder_outputs = F.log_softmax(decoder_outputs,dim=-1)
        return decoder_outputs,decoder_hidden,None
    def forward_step(self,x,hidden):
        x = self.embedding(x)
        x = F.relu(x)
        x,hidden = self.rnn(x,hidden)
        output = self.out(x)
        return output,hidden


decoder = DecoderRNN(hidden_size=5,output_size=10)
target_vector = torch.tensor([[0,1,2,3,4,5,6,7,8,9]])
encoder_outputs,encoder_hiddent = encoder(input_vector)
output,hidden,_=decoder(encoder_outputs,encoder_hiddent,input_vector)
print(output)

